'''
This is a numpy representation of SAM format.
'''
import sys
import numpy as np
import h5py

# define the data type.
ham_100_dt = np.dtype([\
		('QNAME', 'S255'),\
		('FLAG', np.int),\
		('RNAME', 'S100'),\
		('POS', np.int),\
		('MAPQ', np.int),\
		('CIGAR', 'S100'),\
		('RNEXT', 'S100'),\
		('PNEXT', np.int),\
		('TLEN', np.int),\
		('SEQ', 'S100'),\
		('QUAL', 'S100'),\
		('OPTIONAL', 'S50'),\
		])
		
ham_50_dt = np.dtype([\
		('QNAME', 'S255'),\
		('FLAG', np.int),\
		('RNAME', 'S50'),\
		('POS', np.int),\
		('MAPQ', np.int),\
		('CIGAR', 'S50'),\
		('RNEXT', 'S50'),\
		('PNEXT', np.int),\
		('TLEN', np.int),\
		('SEQ', 'S50'),\
		('QUAL', 'S50'),\
		('OPTIONAL','S50'),\
		])
		
ham_var_dt = np.dtype([\
		('QNAME', h5py.new_vlen(str)),\
		('FLAG', np.int),\
		('RNAME', h5py.new_vlen(str)),\
		('POS', np.int),\
		('MAPQ', np.int),\
		('CIGAR', h5py.new_vlen(str)),\
		('RNEXT', h5py.new_vlen(str)),\
		('PNEXT', np.int),\
		('TLEN', np.int),\
		('SEQ', h5py.new_vlen(str)),\
		('QUAL', h5py.new_vlen(str)),\
		('OPTIONAL',h5py.new_vlen(str)),\
		])
		
vlen_dt = h5py.new_vlen(str)
ham_split_dt = np.dtype([\
		('FLAG', np.int),\
		('POS', np.int),\
		('MAPQ', np.int),\
		('PNEXT', np.int),\
		('TLEN', np.int),\
		])
		
initial_size = 1000000
int_chunk_size = 1000000
str_chunk_size = 10000
ham_dt = ham_split_dt

# create a numpy array to fill.
def create_array(self, sz):
	# Declare empty numpy array and return.
	return np.zeros(sz, dtype=self.ham_dt)

class HamWrite(object):
	
	# initializes h5 file and data set.
	def __init__(self, ham_file, init_size=False, ics=False, scs=False):
		# See if initial size, or any of the chunk sizes were set.
		if init_size != False:
			self.initial_size = init_size
		else:
			self.initial_size = initial_size
			
		if ics != False:
			self.int_chunk_size = ics
		else:
			self.int_chunk_size = int_chunk_size
		
		if scs != False:
			self.str_chunk_size = scs
		else:
			self.str_chunk_size = str_chunk_size
		
		# Open h5 file.
		self.h5_out = h5py.File(ham_file, 'w')
		
		# set data type.
		self.ham_dt = ham_dt
		
		# Create h5py object.
		self.data = self.h5_out.create_dataset('nums', (initial_size,), self.ham_dt, maxshape=(None,), compression='gzip', chunks=(self.int_chunk_size,))
		
		# Create string objects.
		self.qnames = self.h5_out.create_dataset('QNAMES', (initial_size,), vlen_dt, maxshape=(None,), compression='gzip', chunks=(self.str_chunk_size,))
		self.rnames = self.h5_out.create_dataset('RNAMES', (initial_size,), vlen_dt, maxshape=(None,), compression='gzip', chunks=(self.str_chunk_size,))
		self.cigars = self.h5_out.create_dataset('CIGARS', (initial_size,), vlen_dt, maxshape=(None,), compression='gzip', chunks=(self.str_chunk_size,))
		self.rnexts = self.h5_out.create_dataset('RNEXTS', (initial_size,), vlen_dt, maxshape=(None,), compression='gzip', chunks=(self.str_chunk_size,))
		self.seqs = self.h5_out.create_dataset('SEQS', (initial_size,), vlen_dt, maxshape=(None,), compression='gzip', chunks=(self.str_chunk_size,))
		self.quals = self.h5_out.create_dataset('QUALS', (initial_size,), vlen_dt, maxshape=(None,), compression='gzip', chunks=(self.str_chunk_size,))
		self.optionals = self.h5_out.create_dataset('OPTIONALS', (initial_size,), vlen_dt, maxshape=(None,), compression='gzip', chunks=(self.str_chunk_size,))
		
		# Setup attributes.
		self.active_idx = 0
		self.max_idx = initial_size


	# closes h5 file.
	def close(self):
		# shrink if necessary.
		if self.max_idx > self.active_idx:
			# Check if we wont seg fault.
			if \
				self.active_idx > self.int_chunk_size \
				and \
				self.active_idx > self.str_chunk_size:

				self.data.resize((self.active_idx,))
			
		# save meta data.
		self.data.attrs['active_idx'] = self.active_idx
		self.data.attrs['sort_type'] = "none"
			
		# close file pointer.
		self.h5_out.close()		

	# save chunk to disk.
	def save_array(self, a, qnames, rnames, cigars, rnexts, seqs, quals, optionals):
		# increase size of h5 dataset.
		new_ptr = self.active_idx + a.size
		if new_ptr > self.max_idx:
			# grow h5 datasets.
			self.data.resize((new_ptr,))
			self.qnames.resize((new_ptr,))
			self.rnames.resize((new_ptr,))
			self.cigars.resize((new_ptr,))
			self.rnexts.resize((new_ptr,))
			self.seqs.resize((new_ptr,))
			self.quals.resize((new_ptr,))
			self.optionals.resize((new_ptr,))
			
			# increase pointer.
			self.max_idx = new_ptr
			
		# save data to disk.
		self.data[self.active_idx:new_ptr] = a
		
		self.qnames[self.active_idx:new_ptr] = qnames
		self.rnames[self.active_idx:new_ptr] = rnames
		self.cigars[self.active_idx:new_ptr] = cigars
		self.rnexts[self.active_idx:new_ptr] = rnexts
		self.seqs[self.active_idx:new_ptr] = seqs
		self.quals[self.active_idx:new_ptr] = quals
		self.optionals[self.active_idx:new_ptr] = optionals
		
		# update final pointer.
		self.active_idx = new_ptr

		
	
class HamRead(object):
	
	# initializes h5 file.
	def __init__(self, ham_file):
		# Open h5 file.
		self.h5_out = h5py.File(ham_file, 'r')

	# closes h5 file.
	def close(self):
		self.h5_out.close()		
		
	# get dataset.
	def get_array(self):
		return self.h5_out['ham_arr']
		
	
